# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import cc.coin1.marketsim.pymarketsim as native_marketsim
from coin.proto import coin_order_enums_pb2
from coin.strategy.marketsim.core import spec_pb2, sim_stat


class MarketSim:
  def __init__(self, sim_spec=None, default_machine=None):
    self._sim_spec = sim_spec
    if self._sim_spec is None:
      self._sim_spec = spec_pb2.MarketSimSpecProto(
          latency=10**9,  # 1 sec
          feed_cut_time_allowance=60 * 10**9,  # 1 min
          clear_on_feed_cut=True)

    self._default_machine = default_machine
    self._products = {}
    self._symbol_to_product_map = {}
    self._execution_callbacks = []

    self._native_sim = native_marketsim._MarketSim()
    self._native_sim.init(self._sim_spec.SerializeToString(), self._on_execution)

  @property
  def default_machine(self):
    return self._default_machine

  def add_product(self, product, product_spec):
    pspec = spec_pb2.ProductSpecProto()
    pspec.CopyFrom(product_spec)
    pspec.fastfeed_machine = pspec.fastfeed_machine or self._default_machine

    assert pspec.symbol == str(product)
    assert pspec.fastfeed_sub_req

    self._products[product] = pspec
    self._symbol_to_product_map[str(product)] = product
    self._native_sim.add_product(pspec.SerializeToString())

  def get_product_spec(self, product):
    if isinstance(product, str):
      product = self._symbol_to_product_map[product]
    return self._products[product]

  def get_product_from_symbol(self, symbol):
    return self._symbol_to_product_map[symbol]

  def submit(self,
             *,
             product,
             side,
             price,
             qty,
             order_type=coin_order_enums_pb2.DAY_LIMIT_ORDER,
             post_only=False,
             client_order_id=None,
             tag=None):
    # TODO(jaewon): sanitize price, qty unit size
    assert qty > 0, qty
    self._native_sim.submit(str(product),
                            side,
                            order_type,
                            price,
                            qty,
                            post_only,
                            client_order_id,
                            tag)

  def cancel(self, *, order_id=None, client_order_id=None):
    assert (order_id is None) != (client_order_id is None)
    self._native_sim.cancel(order_id, client_order_id)

  def cancel_all(self):
    for order in self.get_orders():
      if order.order_status == spec_pb2.ORDER_STATUS_WORKING and not order.cancel_submitted:
        self.cancel(order_id=order.order_id)

  def get_orders(self):
    return self._native_sim.get_orders()

  def get_order(self, *, order_id=None, client_order_id=None):
    assert (order_id is None) != (client_order_id is None)
    if order_id is not None:
      assert isinstance(order_id, int)
      return self._native_sim.get_order(order_id)
    else:
      assert isinstance(client_order_id, str)
      return self._native_sim.get_order(client_order_id)

  def get_positions(self):
    return self._native_sim.get_positions()

  def get_position(self, product):
    return self._native_sim.get_position(str(product))

  def get_order_history(self):
    serialized_orders = self._native_sim.get_order_history()
    orders = []
    for serialized in serialized_orders:
      order = spec_pb2.OrderProto()
      order.ParseFromString(serialized)
      orders.append(order)
    return orders

  def get_execution_history(self):
    serialized_execs = self._native_sim.get_execution_history()
    execs = []
    for serialized in serialized_execs:
      execution = spec_pb2.ExecutionProto()
      execution.ParseFromString(serialized)
      execs.append(execution)
    return execs

  def dump_sim_result(self, filepath, tag=None):
    self._native_sim.dump_sim_result(filepath, tag)

  def generate_sim_stat(self):
    positions = self.get_positions()
    orders = sorted(self.get_order_history() + self.get_orders(), key=lambda x: x.submit_ts)
    executions = self.get_execution_history()
    product_specs = sorted(self._products.values(), key=lambda x: x.symbol)
    return sim_stat.generate_sim_stat(positions, orders, executions, product_specs)

  def add_execution_callback(self, callback):
    self._execution_callbacks.append(callback)

  def _on_execution(self, serialized):
    if not self._execution_callbacks:
      return
    execution = spec_pb2.ExecutionProto()
    execution.ParseFromString(serialized)
    for callback in self._execution_callbacks:
      callback(execution)

  def get_feed_sub_reqs(self, feed_sub_req=None):
    from coin.strategy.mm.subscription import SubscriptionRequest, FeedSubscriptionRequest
    feed_sub_req = feed_sub_req or FeedSubscriptionRequest()
    for req in self._native_sim.fastfeed_sub_reqs:
      assert req.machine == self.default_machine
      sub_req = SubscriptionRequest.from_str(req.sub_req)
      products = [self._symbol_to_product_map[symbol] for symbol in req.symbols]
      feed_sub_req.add_products(products, sub_req)
    return feed_sub_req

  def on_book_reset(self, sub_req, book_builder, feed_cache_key):
    if sub_req is not None:
      assert feed_cache_key.sub_req == sub_req.to_str()
    self._native_sim.on_book_reset(feed_cache_key, book_builder.cc_builder)
